package smartai.examples.face;

import cn.smartjavaai.face.*;
import cn.smartjavaai.face.entity.FaceResult;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.junit.Assert;
import org.junit.Test;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;

/**
 * FaceNet人脸算法模型demo
 * 支持功能：人脸特征提取、人脸比对（1：1）
 * @author dwj
 * @date 2025/4/11
 */
@Slf4j
public class FaceNetDemo {


    /**
     * 人脸比对（1：1）-在线模型
     * 图片参数：图片路径
     * @throws Exception
     */
    @Test
    public void featureComparison(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            //config.setModelPath("/Users/xxx/Documents/develop/face_model/model_ir_se50.pth");
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            //自动裁剪人脸并比对人脸特征
            float similar = faceModel.featureComparison("src/main/resources/kana1.jpg","src/main/resources/kana2.jpg");
            log.info("相似度：{}", similar);
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）- 使用离线模型
     * 图片参数：图片路径
     * @throws Exception
     */
    @Test
    public void featureComparisonOffline(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            config.setModelPath("/Users/xxx/Documents/develop/face_model/face_feature.pt");
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            //自动裁剪人脸并比对人脸特征
            float similar = faceModel.featureComparison("src/main/resources/kana1.jpg","src/main/resources/kana2.jpg");
            log.info("相似度：{}", similar);
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）
     * 图片参数：BufferedImage
     */
    @Test
    public void featureComparison2(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            BufferedImage image1 = ImageIO.read(new File(Paths.get("src/main/resources/kana1.jpg").toAbsolutePath().toString()));
            BufferedImage image2 = ImageIO.read(new File(Paths.get("src/main/resources/kana2.jpg").toAbsolutePath().toString()));
            //自动裁剪人脸并比对人脸特征
            float similar = faceModel.featureComparison(image1,image2);
            log.info("相似度：{}", similar);
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）
     * 图片参数：byte[]
     */
    @Test
    public void featureComparison3(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            byte[] imageData1 = Files.readAllBytes(Paths.get("src/main/resources/kana1.jpg"));
            byte[] imageData2 = Files.readAllBytes(Paths.get("src/main/resources/kana2.jpg"));
            //自动裁剪人脸并比对人脸特征
            float similar = faceModel.featureComparison(imageData1,imageData2);
            log.info("相似度：{}", similar);
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）
     * 先特征提取，后比对人脸特征
     * 提取人脸特征图片参数：图片路径
     */
    @Test
    public void featureExtractionAndCompare(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            //提取图像中最大人脸的特征
            float[] feature1 = faceModel.featureExtraction("src/main/resources/kana1.jpg");
            float[] feature2 = faceModel.featureExtraction("src/main/resources/kana2.jpg");
            if(feature1 != null && feature2 != null){
                float similar = faceModel.calculSimilar(feature1, feature2);
                log.info("相似度：{}", similar);
            }else{
                log.warn("人脸特征提取失败");
            }
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）
     * 先特征提取，后比对人脸特征
     * 提取人脸特征图片参数：BufferedImage
     */
    @Test
    public void featureExtractionAndCompare2(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            BufferedImage image1 = ImageIO.read(new File(Paths.get("src/main/resources/kana1.jpg").toAbsolutePath().toString()));
            BufferedImage image2 = ImageIO.read(new File(Paths.get("src/main/resources/kana2.jpg").toAbsolutePath().toString()));
            //提取图像中最大人脸的特征
            float[] feature1 = faceModel.featureExtraction(image1);
            float[] feature2 = faceModel.featureExtraction(image2);
            if(feature1 != null && feature2 != null){
                float similar = faceModel.calculSimilar(feature1, feature2);
                log.info("相似度：{}", similar);
            }else{
                log.warn("人脸特征提取失败");
            }
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 人脸比对（1：1）
     * 先特征提取，后比对人脸特征
     * 提取人脸特征图片参数：byte[]
     */
    @Test
    public void featureExtractionAndCompare3(){
        try {
            FaceModelConfig config = new FaceModelConfig();
            config.setModelEnum(FaceModelEnum.FACENET_FEATURE_EXTRACTION);//人脸模型
            FaceModel faceModel = FaceModelFactory.getInstance().getModel(config);
            byte[] imageData1 = Files.readAllBytes(Paths.get("src/main/resources/kana1.jpg"));
            byte[] imageData2 = Files.readAllBytes(Paths.get("src/main/resources/kana2.jpg"));
            //提取图像中最大人脸的特征
            float[] feature1 = faceModel.featureExtraction(imageData1);
            float[] feature2 = faceModel.featureExtraction(imageData2);
            if(feature1 != null && feature2 != null){
                float similar = faceModel.calculSimilar(feature1, feature2);
                log.info("相似度：{}", similar);
            }else{
                log.warn("人脸特征提取失败");
            }
        }
        catch (Exception e){
            e.printStackTrace();
        }
    }


}
